package com.example.socket.websocket;

import com.example.socket.ServletRequest;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * @Classname WebSocketServer
 * @Description TODO
 * @Date 2023/4/27 17:45
 * @Created by XiongXiong
 * @Author: X.I.O
 */
public class WebSocketServer implements Runnable{

    private ServletRequest request;

    private OutputStream outputStream;

    @Override
    public void run() {
        try {
            while (true) {
                startWebSocket(request,outputStream);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public  void startWebSocket(ServletRequest request, OutputStream out) throws Exception {

        // 读取请求头信息
        Map<String, String> headers = request.getHeards();
        if (!headers.containsKey("Upgrade") || !headers.get("Upgrade").equals("websocket")) {
            System.out.println("Invalid WebSocket request.");
            return;
        }
        InputStream inputStream = request.getInputStream();

        // 获取 WebSocket key，用于生成响应头中的 Sec-WebSocket-Accept 字段
        String key = headers.get("Sec-WebSocket-Key");
        if (key == null) {
            return;
        }

        // 生成 Sec-WebSocket-Accept 响应头字段值
        String accept = generateAccept(key);
        // 发送 HTTP 响应头，升级为 WebSocket 连接
        writeLine(outputStream, "HTTP/1.1 101 Switching Protocols");
        writeLine(outputStream, "Upgrade: websocket");
        writeLine(outputStream, "Connection: Upgrade");
        writeLine(outputStream, "Sec-WebSocket-Accept: " + accept);
        writeLine(outputStream, "");
        // 处理 WebSocket 连接
        handleWebSocket(inputStream, outputStream);
    }

    private  Map<String, String> readHeaders(InputStream inputStream) throws Exception {
        Map<String, String> headers = new HashMap<>();
        while (true) {
            String line = readLine(inputStream);
            if (line.isEmpty()) {
                break;
            }
            int index = line.indexOf(":");
            if (index > 0) {
                headers.put(line.substring(0, index).trim(), line.substring(index + 1).trim());
            }
        }
        return headers;
    }

    private  String readLine(InputStream inputStream) throws Exception {
        StringBuilder sb = new StringBuilder();
        int c;
        while ((c = inputStream.read()) != -1) {
            if (c == '\r') {
                continue;
            } else if (c == '\n') {
                break;
            } else {
                sb.append((char) c);
            }
        }
        return sb.toString();
    }

    private void writeLine(OutputStream outputStream, String line) throws Exception {
        outputStream.write(line.getBytes());
        outputStream.write("\r\n".getBytes());
    }

    private  String generateAccept(String key) throws NoSuchAlgorithmException {
        final String magicString = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
        MessageDigest sha1 = MessageDigest.getInstance("SHA-1");
        String input = key + magicString;
        byte[] hash = sha1.digest(input.getBytes());
        return Base64.getEncoder().encodeToString(hash);
    }

    private  void handleWebSocket(InputStream inputStream, OutputStream outputStream) throws Exception {
        byte[] buffer = new byte[1024];
        while (true) {
            int len = inputStream.read(buffer);
            if (len < 0) {
                break;
            }
            // 处理 WebSocket 消息并回显
            String message = decodeMessage(buffer, len);
            System.out.println("Received message: " + message);
            send(outputStream, "Echo  后台: " + message);
        }
    }

    private  String decodeMessage(byte[] buffer, int length) {
        byte b1 = buffer[0];
        boolean fin = (b1 & 0x80) != 0;
        int opcode = b1 & 0x0F;
        byte b2 = buffer[1];
        boolean masked = (b2 & 0x80) != 0;
        int payloadLen = b2 & 0x7F;

        int offset = 2;
        if (payloadLen == 126) {
            payloadLen = ((buffer[2] & 0xFF) << 8) | (buffer[3] & 0xFF);
            offset = 4;
        } else if (payloadLen == 127) {
            payloadLen = ((buffer[2] & 0xFF) << 56)
                    | ((buffer[3] & 0xFF) << 48)
                    | ((buffer[4] & 0xFF) << 40)
                    | ((buffer[5] & 0xFF) << 32)
                    | ((buffer[6] & 0xFF) << 24)
                    | ((buffer[7] & 0xFF) << 16)
                    | ((buffer[8] & 0xFF) << 8)
                    | (buffer[9] & 0xFF);
            offset = 10;
        }

        byte[] maskingKey = null;
        if (masked) {
            maskingKey = new byte[]{buffer[offset], buffer[offset + 1], buffer[offset + 2], buffer[offset + 3]};
            offset += 4;
        }

        byte[] payload = new byte[payloadLen];
        System.arraycopy(buffer, offset, payload, 0, payloadLen);

        if (masked) {
            for (int i = 0; i < payloadLen; i++) {
                payload[i] ^= maskingKey[i % 4];
            }
        }

        return new String(payload);
    }

    private  void send(OutputStream outputStream, String message) throws IOException {
        // 将消息编码成二进制帧，并写入输出流
        byte[] bytes = encodeMessage(message);
        outputStream.write(bytes);
        outputStream.flush();
    }


    private  byte[] encodeMessage(String message) {
        byte[] messageBytes = message.getBytes();
        int payloadLen = messageBytes.length;
        byte[] header;
        if (payloadLen <= 125) {
            header = new byte[2];
            header[1] = (byte) payloadLen;
        } else if (payloadLen <= 65535) {
            header = new byte[4];
            header[1] = (byte) 126;
            header[2] = (byte) ((payloadLen >> 8) & 0xFF);
            header[3] = (byte) (payloadLen & 0xFF);
        } else {
            header = new byte[10];
            header[1] = (byte) 127;
            header[2] = (byte) ((payloadLen >> 56) & 0xFF);
            header[3] = (byte) ((payloadLen >> 48) & 0xFF);
            header[4] = (byte) ((payloadLen >> 40) & 0xFF);
            header[5] = (byte) ((payloadLen >> 32) & 0xFF);
            header[6] = (byte) ((payloadLen >> 24) & 0xFF);
            header[7] = (byte) ((payloadLen >> 16) & 0xFF);
            header[8] = (byte) ((payloadLen >> 8) & 0xFF);
            header[9] = (byte) (payloadLen & 0xFF);
        }

        header[0] = (byte) 0x81;

        byte[] frame = new byte[header.length + messageBytes.length];
        System.arraycopy(header, 0, frame, 0, header.length);
        System.arraycopy(messageBytes, 0, frame, header.length, messageBytes.length);

        return frame;
    }

    public WebSocketServer(ServletRequest request, OutputStream outputStream) {
        this.request = request;
        this.outputStream = outputStream;
    }
}
